#!/usr/bin/env python
from __future__ import print_function
from pyexotica.transformations import euler_from_quaternion

def parseSceneFile(file_name, debug=False):
    data = []
    with open(file_name, 'rU') as fin:
        for line in fin:
            line = line.strip()
            if line:
                data.append(line)

    objects = {}

    # Get name of scene
    if data[0][0] != "(" or data[0][-2:] != ")+":
        raise RuntimeError("Invalid scene file - does not start with (SceneName)+.")
    sceneName = data[0].replace("(", "").replace(")+", "")

    if debug:
        print("Parsing Scene:", sceneName)
    idx = 1
    while idx < len(data):
        if data[idx][0] == ".":
            if debug:
                print("Done")
            break
        elif data[idx][0] == "*":
            objectName = data[idx].replace("* ", "")
            idx += 1
            numberOfElements = int(data[idx])
            if debug:
                print("New object:", objectName,
                      "with", numberOfElements, "elements")
            components = []
            for i in range(numberOfElements):
                if data[idx + 1] == "box":
                    size = [float(e) for e in data[idx + 2].split(" ")]
                    position = [float(e) for e in data[idx + 3].split(" ")]
                    orientation = [float(e) for e in data[idx + 4].split(" ")]
                    color = [float(e) for e in data[idx + 5].split(" ")]
                    if debug:
                        print("Box #" + str(i) + ": Size=", size,
                              "Position=", position,
                              "Orientation=", orientation,
                              "Color=", color)
                    components.append({'type': 'box',
                                       'size': size,
                                       'position': position,
                                       'orientation': orientation,
                                       'color': color})
                    idx += 5
                else:
                    raise Exception("Not supported", data[idx + 1])
            objects.update({objectName: components})
        else:
            idx += 1

    return objects


def writeBoxGeometryToSdfFile(fileSdf, geometry, objectName, objectNumber):
    position = geometry['position']
    orientation = geometry['orientation']
    rpy = euler_from_quaternion(
        [orientation[3], orientation[0], orientation[1], orientation[2]])
    size = geometry['size']
    color = geometry['color']
    fileSdf.write("            <link name='%s'>\n"
                  "                <pose>%f %f %f %f %f %f</pose>\n"
                  "                <collision name='collision'>\n"
                  "                    <geometry>\n"
                  "                        <box>\n"
                  "                            <size>%f %f %f</size>\n"
                  "                        </box>\n"
                  "                    </geometry>\n"
                  "                </collision>\n"
                  "                <visual name='visual'>\n"
                  "                    <geometry>\n"
                  "                        <box>\n"
                  "                            <size>%f %f %f</size>\n"
                  "                        </box>\n"
                  "                    </geometry>\n"
                  "                    <material>\n"
                  "                        <ambient>%f %f %f</ambient>\n"
                  "                        <diffuse>%f %f %f %f</diffuse>\n"
                  "                    </material>\n"
                  "                </visual>\n"
                  "            </link>\n"
                  % (objectName + "_" + str(objectNumber),
                     position[0], position[1], position[2],
                     rpy[0], rpy[1], rpy[2],
                     size[0], size[1], size[2], size[0], size[1], size[2],
                     color[0], color[1], color[2],
                     color[0], color[1], color[2], color[3]))


def createSdfFromObjects(sdf_file_name, objects, pose):
    fileSdf = open(sdf_file_name, "w")

    fileSdf.write("<sdf version='1.4'>\n"
                  "    <world name='demo'>\n")

    for objectName in objects:
        fileSdf.write("        <model name='%s'>\n"
                      "            <pose>%f %f %f %f %f %f</pose>\n"
                      % (objectName,
                         pose[0], pose[1], pose[2], pose[3], pose[4], pose[5]))
        objectCount = 0
        for geometry in objects[objectName]:
            if geometry['type'] == 'box':
                writeBoxGeometryToSdfFile(
                    fileSdf, geometry, objectName, objectCount)
            else:
                raise Exception("Not supported: ", geometry['type'])
            objectCount += 1
        fileSdf.write("            <static>1</static>\n"
                      "        </model>\n")

    fileSdf.write("    </world>\n"
                  "</sdf>\n")
    fileSdf.close()


if __name__ == "__main__":
    import os
    import sys
    import numpy as np

    if len(sys.argv) != 3:
        print("Two arguments required: [scene input file] [sdf output file]")
    scene_file = os.path.abspath(sys.argv[1])
    sdf_file = os.path.abspath(sys.argv[2])

    # Check that scene exists
    if not os.path.exists(scene_file):
        raise RuntimeError("Scene file does not exist!")
    if os.path.exists(sdf_file):
        print("WARNING: SDF exists, will be overwritten.")

    # TODO THIS IS HARD-CODED, MAKE PARAMETER
    pose = [0.] * 6 #[0.8, 0, 0, 0, 0, np.pi/2]
    objects = parseSceneFile(scene_file)
    createSdfFromObjects(sdf_file, objects, pose)
    print(">>> DONE <<<")
